# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
import matplotlib.finance as mpf
import os
import time
import datetime
from matplotlib.dates import date2num,num2date
from matplotlib.font_manager import FontProperties
# from matplotlib.finance import date2num,num2date
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
from matplotlib.ticker import Formatter
from matplotlib.gridspec import GridSpec
from matplotlib.widgets import Button
from multiprocessing import Process

datapath='d:/tushare_datas'
font = FontProperties(fname='c:/simsun.ttf',size=12,weight=1)

fig=plt.figure()
gs=GridSpec(3,1)
ax1=plt.subplot(gs[:-1,0])
ax2=plt.subplot(gs[-1,0],sharex=ax1)
fig.subplots_adjust(bottom=0.2)

class MyFormatter(Formatter):
	def __init__(self,dates,fmt='%Y-%m-%d'):
		self.dates=dates
		self.fmt=fmt

	def __call__(self,x,pos=0):
		'Return the label for time x at position pos'
		ind=int(np.round(x))
		if ind>=len(self.dates) or ind<0:
			return ''
		return self.dates[ind].strftime(self.fmt)	

class MyIndex(object):
	ind=0
	cid=0
	rid=0
	mid=0
	pid=0
	isDrawLine=False
	Coords1x=0
	Coords1y=0
	"""docstring for MyIndex"""
	def __init__(self, codes, ax):
		# super(MyIndex, self).__init__()
		self.codes = codes
		self.isDrawLine=False
		self.ax = ax

	def OnClick(self,event):   
	    if event.button ==1:  
	        self.Coords1x = event.xdata  
	        self.Coords1y = event.ydata  

	def OnMouseMotion(self,event):  
	    if event.button == 1:  
	        # Coords2x = event.xdata  
	        # Coords2y = event.ydata  
	        x1 = [self.Coords1x,event.xdata]  
	        y1 = [self.Coords1y,event.ydata]          
	        # ax = plt.gca()  
	        lines = self.ax.plot(x1,y1,picker = 20)  
	        # print(lines)
	        self.ax.figure.canvas.draw()  
	        #删除之前的线条，进行更新  
	        l = lines.pop(0)  
	        l.remove()  
	        del l  

	def OnRelease(self,event):
	    # x1 = [Coords1x,event.xdata]  
	    # y1 = [Coords1y,event.ydata]  
	    if self.Coords1x==event.ydata:
	    	a=0
	    else:
	        a=(self.Coords1y - event.ydata)/(self.Coords1x - event.xdata)
	    b=self.Coords1y-a*self.Coords1x
	    l,r=self.ax.get_xlim()
	    bottom,top=self.ax.get_ylim()
	    x1 = np.arange(l,r)  
	    y1 = a*x1+b
	    # ax1=plt.gca()
	    self.ax.plot(x1,y1,picker=20)
	    self.ax.figure.canvas.draw()

	def OnPick(self,event):
	    line=event.artist
	    line.remove()
	    del line

	def next(self,event):
		self.ind +=1
		if self.ind<len(self.codes):
			stocks_view2(self.codes.loc[self.ind])

	def prev(self,event):
		self.ind -=1
		if self.ind>=0:
			stocks_view2(self.codes.loc[self.ind])

	def drawline(self,event):
		if self.isDrawLine:
			event.canvas.figure.patch.set_facecolor('red')
			self.cid=fig.canvas.mpl_connect('button_press_event',self.OnClick)  
			self.rid=fig.canvas.mpl_connect('button_release_event',self.OnRelease)
			self.mid=fig.canvas.mpl_connect('motion_notify_event',self.OnMouseMotion) 
			self.pid=fig.canvas.mpl_connect('pick_event',self.OnPick)
		else:
			event.canvas.figure.patch.set_facecolor('yellow')
			fig.canvas.mpl_disconnect(self.cid)
			fig.canvas.mpl_disconnect(self.rid)
			fig.canvas.mpl_disconnect(self.mid)
			fig.canvas.mpl_disconnect(self.pid)
		self.isDrawLine=not self.isDrawLine

def get_stock_data():
	stocks=ts.get_stock_basics().index.tolist()
	for code in stocks:
		print('Starting building %s.csv' % code)
		df=ts.get_hist_data(code)
		df.to_csv('%s/%s.csv' % (datapath,code))

def get_stockpool_names():
	stocks=ts.get_stock_basics()
	with open('stockpool.txt','r') as f:
		s=f.read().split(',')
	return stocks.loc[stocks.index.isin(s),'name'].reset_index()

def update_stock_data():
	stocks=get_stock_list(datapath)
	for code in stocks:
		# lines=[]
		# with open('%s/%s.csv' % (datapath,code)) as f:
		# 	lines=f.readlines()
		# h,d=lines[0],lines[1:]
		print('[ %s ] update starting' % code)
		file='%s/%s.csv' % (datapath,code)
		olddata=pd.DataFrame.from_csv(file)
		start_time=olddata.index[0]+datetime.timedelta(days=1)
		newdata=ts.get_hist_data(code,start=start_time.date().strftime('%Y-%m-%d'))
		newdata.index=pd.to_datetime(newdata.index)
		data=pd.concat([newdata,olddata])
		data.to_csv(file)


def datestr2num(s):
	return datetime.datetime.strptime(s.decode('ascii'),'%Y-%m-%d').timestamp()

def date_to_num(date):
	return date2num(datetime.datetime.strptime(date,'%Y-%m-%d'))

def get_ma(data,N=5):
	ma=[]
	M=len(data)-N
	for i in range(M):
		ma.append(np.mean(data[i:i+N]))
	return ma

def get_macd_data(data,short=12,long1=26,mid=9):
	data['sema']=data['close'].ewm(span=short).mean()
	data['lema']=data['close'].ewm(span=long1).mean()
	data.fillna(0,inplace=True)
	data['data_dif']=data['sema']-data['lema']
	data['data_dea']=data['data_dif'].ewm(span=mid).mean()
	data['data_macd']=2*(data['data_dif']-data['data_dea'])
	data.fillna(0,inplace=True)
	return data[['data_dif','data_dea','data_macd']]

def get_kdj_data(data,N=9,M=2):
	low_list=data['low'].rolling(N).min()
	low_list.fillna(value=data['low'].expanding().min(),inplace=True)
	high_list=data['high'].rolling(N).max()
	high_list.fillna(value=data['high'].expanding().max(),inplace=True)
	rsv=(data['close']-low_list)/(high_list-low_list)*100
	data['kdj_k']=rsv.ewm(com=M).mean()
	data['kdj_d']=data['kdj_k'].ewm(com=M).mean()
	data['kdj_j']=3*data['kdj_k']-2*data['kdj_d']
	data.fillna(0,inplace=True)
	return data[['kdj_k','kdj_d','kdj_j']]

def simple_ma(data,window_short=5,window_middle=20,window_long=60):
	data['ma_short']=data['close'].rolling(window_short).mean()
	data['ma_middle']=data['close'].rolling(window_middle).mean()
	data['ma_long']=data['close'].rolling(window_long).mean()
	data['ma_short'].fillna(0,inplace=True)
	data['ma_middle'].fillna(0,inplace=True)
	data['ma_long'].fillna(0,inplace=True)
	data.ix[(data['ma_short'].shift(1)>data['ma_long'].shift(1))&(data['open']<data['close'].shift(1)*1.097),'position']=1
	data.ix[(data['ma_short'].shift(1)<data['ma_long'].shift(1))&(data['open']>data['close'].shift(1)*0.903),'position']=0
	data['position'].fillna(method='ffill',inplace=True)
	data['position'].fillna(0,inplace=True)
	return data[['ma_short','ma_middle','ma_long','position']]

def get_stock_list(path):
	codes=[]
	fileList=os.listdir(path)
	for f in fileList:
		codes.append(os.path.splitext(f)[0])
	return codes

def cacl_ma_fit1d(t,ma,d):
	t1=t[:d]
	ma1=ma[:d]
	f1=np.polyfit(t1,ma1,1)
	return f1[0]

def stocks_view2(code):
	ax1.cla()
	ax2.cla()
	data=pd.read_csv('%s/%s.csv' % (datapath,code['code']))

	data.sort_index(ascending=False,inplace=True)
	simple_ma(data)
	data.sort_index(inplace=True)
	L=180
	data=data[:L]

	numtime=data['date'].apply(date_to_num)
	numtime.sort_index(ascending=False,inplace=True)
	# print(numtime)
	data['date']=pd.Series(np.arange(L,0,-1))
	matdata=data.ix[:,['date','open','close','high','low','volume']]

	matdata.rename(columns={'date':'time'},inplace=True)
	matdata=matdata.as_matrix()

	xs=[num2date(tm).date() for tm in numtime]
	formatter=MyFormatter(xs)
	fig.suptitle('%s %s' %(code['code'],code['name']),fontproperties=font)
	ax1.xaxis.set_major_formatter(formatter)

	mpf.candlestick_ochl(ax1,matdata,width=0.6,colorup='r',colordown='g')
	ax1.plot(matdata[:,0],data['ma_short'],label='ma5')
	ax1.plot(matdata[:,0],data['ma_middle'],label='ma20')
	ax1.plot(matdata[:,0],data['ma_long'],label='ma60')
	ax1.grid(True)
	ax1.set_xlim(L-90,L)
	
	# ax1.xaxis_date()
	ax2.bar(matdata[:,0],matdata[:,5],width=0.5)
	ax2.grid(True)
	plt.setp(ax1.get_xticklabels(),visible=False)
	# fig.autofmt_xdate()	

def check_stock_pool(pstocks):

	codes=get_stock_list(datapath)

	addCodes=[]

	for index,c in enumerate(codes):
		print("start check",c)
		cdata=np.loadtxt(datapath+os.sep+c+'.csv',delimiter=',',usecols=(0,3),converters={0:datestr2num},skiprows=1,unpack=True)
		if not isinstance(cdata[0],np.ndarray) or len(cdata[0])<60:
			continue
		if datetime.datetime.fromtimestamp(cdata[0][0]).strftime('%Y-%m-%d')!=datetime.datetime.now().strftime('%Y-%m-%d'):
			if c in pstocks:
				print('[ %s ] check out' % c)
				break
				pstocks.remove(c)
			continue

		m60=np.mean(cdata[1][0:60])
		m5=np.mean(cdata[1][0:5])
		# print(c,' 60 np 5 ',m60,m5)
		if m5>m60 :
			if not(c in pstocks):
				print('[ %s ] check in' % c)
				addCodes.append(c)

		else:
			if  c in pstocks:
				print('[ %s ] check out' % c)
				pstocks.remove(c)

	pstocks.extend(addCodes)
	with open('stockpool.txt','w') as f:
		# np.savetxt(f,pstocks,delimiter=',',header='触发日期,代码',fmt='%s',newline='\r\n',encoding='utf8')
		f.write(','.join(pstocks))
	return pstocks

def draw_stocks(ax,ddata,odata,hdata,cdata,ldata,upcolor='r',downcolor='g'):
	upboxes=[]
	downboxes=[]
	for d,o,h,c,l in zip(ddata,odata,hdata,cdata,ldata):
		rect=Rectangle((d-0.4,min(o,c)),0.8,abs(o-c))
		rect2=Rectangle((d-0.01,l),0.02,h-l)
		if c>=o :
			upboxes.append(rect)
			upboxes.append(rect2)
		else:
			downboxes.append(rect)
			downboxes.append(rect2)
	upc=PatchCollection(upboxes,facecolor=upcolor,edgecolor=upcolor)
	dpc=PatchCollection(downboxes,facecolor=downcolor,edgecolor=downcolor)
	ax.add_collection(upc)
	ax.add_collection(dpc)

def run_proc(name):
	# print('Run child process %s (%s)' % (name,os.getpid()))
	pstocks=[]
	with open('stockpool.txt',encoding='utf8') as f:
		# pstocks=np.loadtxt(f,dtype=str,delimiter=',',usecols=(0),skiprows=1,unpack=True)
		str=f.read()
		if len(str)!=0:
			pstocks=str.split(',')

	
	# stocks_view(pstocks[0])
	stockname=get_stockpool_names()

	callback=MyIndex(stockname,ax1)
	# stocks_view2(pstocks[0])
	stocks_view2(stockname.loc[0])
	
	axnext=plt.axes([0.81,0.05,0.1,0.075])
	axprev=plt.axes([0.7,0.05,0.1,0.075])
	axline=plt.axes([0.5,0.05,0.1,0.075])
	# axline.patch.set_facecolor('blue')
	bnext=Button(axnext,'Next')
	bprev=Button(axprev,'Prev')
	bline=Button(axline,'DrawLine')
	bnext.on_clicked(callback.next)
	bprev.on_clicked(callback.prev)
	bline.on_clicked(callback.drawline)

	plt.show()	

# stocks_view('000001')

if __name__ == '__main__':
	p='''
	#########选择菜单，输入相应数字：########
	###
	###  1 -- 检查stockpool
	###  2 -- 显示stocklist
	###  3 -- 下载数据
	###  4 -- 更新数据
	###  2 -- 显示stocklist
	###  0 -- 退出
	###
	#########################################
	'''
	while(True):
		print(p)
		num=input("### 请选择：")
		# dictName={1:check_stock_pool,2:show_stocks,}
		if num=='1':
			check_stock_pool([])
		elif num=='2':
			p=Process(target=run_proc,args=('test',))
			p.start()
			# p.join()
			# print('Child process end.')
		elif num=='3':
			get_stock_data()
		elif num=='4':
			update_stock_data()
		elif num=='0':
			
			break
		else:
			break




